Passer de mode immédiat PyTorch à Triton exige un changement de perspective : au lieu de considérer les tenseurs comme des objets monolithiques, il faut les voir comme des collections de blocs discrets et gérables blocs ou tuiles.
1. Tenseurs PyTorch vs. tenseurs Triton
Il est essentiel de distinguer tenseurs Triton de tenseurs PyTorch. Un tenseur PyTorch est un objet Python côté hôte enveloppant la forme, le type de données (dtype), l'appareil (device), les décalages (strides) et les métadonnées de stockage. En revanche, Triton travaille avec les pointeurs vers les données brutes dans des blocs mémoire spécifiques, permettant une optimisation beaucoup plus fine au niveau du matériel.
2. Le goulot d'étranglement de l'exécution immédiate
Dans l'exécution standard immédiate, chaque opération (par exemple, addition puis ReLU) nécessite un lancement séparé de noyau et un aller-retour en mémoire globale. C'est le principal goulot d'étranglement dans le calcul moderne sur GPU. Triton contourne ce problème en fusionnant des opérations au sein d'un seul noyau qui traite des blocs de données (par exemple, 128, 256 ou 512 éléments) directement en mémoire intégrée.
3. Le paradigme basé sur les blocs
Plutôt que de penser au niveau scalaire comme dans les threads CUDA, Triton utilise SPMD (programme unique, multiples données) au niveau du bloc. Vous écrivez un seul noyau, et Triton lance plusieurs instances à travers une grille. Chaque instance utilise son program_id pour calculer quel "morceau" de mémoire lui appartient.
4. Configuration de l'environnement
Pour commencer, installez Triton dans un environnement propre (en utilisant Conda ou venv) pour éviter tout conflit de dépendances avec les kits CUDA existants : pip install triton.